#include "SimulationManager.h"
#include "FluidField.h"
#include "FluidSolver.h"
#include "Field.h"
#include "haarwavelet.h"
#include "HaarSelector.h"

#include <glm/glm.hpp>

#define WINDOW_W 600
#define WINDOW_H 600

unsigned int fluidDimX = 128, fluidDimY = 128;
float appliedForce = 50.0f;
float fluidVisc =  0.001f;
float fieldBorder = 0.0f; //The amount of border on the edge.
float ts = 0.01f;
int fkernel = 5; // applied force will be in a fkernel*2 X fkernel*2 region
int skernel = 4;
float vcEnergy = 0.1f;
bool doDecomp= false;
bool doDrawSelectedCells = true;
int decompLev = 0;

HaarWavelet haarWv;

int selectedCellX, selectedCellY;

bool doDrawVectorField = false;

HaarSelector haarSelector(fluidDimX, fluidDimY, haarWv.GetDecomp()->numLevs);

Field<float> curl;
Field<glm::vec2> curlGrad;
Field<glm::vec2> curlField;
Field<glm::vec2> diffField;
Field<glm::vec2> waveField;
Field<float> swaveField;
Field<glm::vec2> testField;


//A global simulation manager.
//This takes care of most boiler-plate code.
//If you only want to draw something, you only need to do three things:
//1) Create a display func with the function signerature
//	void funcname()
//2) In main, call 
//	sm.SetCallbackDisplay(funcname)
//  where funcname is the name of your draw function.
//3) In main, after the above function call, call
//	sm.Run();
//  This will start running the simulation.
// 
//  For a list of the callbacks the user can set, look at the "SetCallback..." functions 
//  in the SimulationManager header file.
//  Also, the GetCamera() method will get the camera and allow you to call methods which move
//	around the scene or rotate, etc.
SimulationManager sm(WINDOW_W, WINDOW_H);
Field<glm::vec2> * field;
FluidSolver * fs;


void InitFrame(){
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();


	glOrtho(0-fieldBorder, 1+fieldBorder, 0-fieldBorder, 1+fieldBorder, 0.1, 1000.0);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	gluLookAt(0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0);

	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
}

void DrawScalarField( Field<float> & f, int dimx, int dimy ){

	glBegin(GL_QUADS);
	for (int i = 1; i < dimx-1; ++i){
		for (int j=1; j< dimy-1;++j){
			float val;
			/*if (std::fabs(val) > 0)*/{
				f.Sample( (float)(i)/(dimx-1), (float)(j)/(dimy-1), val );
				glColor3f(val, val, val);
				glVertex2f((float)i/dimx, (float)j/dimy);
				glVertex2f((float)(i+1)/dimx, (float)j/dimy);
				glVertex2f((float)(i+1)/dimx, (float)(j+1)/dimy);
				glVertex2f((float)i/dimx, (float)(j+1)/dimy);
			}
		}
	}
	glEnd();
}

void DrawVectorField(Field<glm::vec2> & f, int dimx, int dimy, float mulConst = 1.0f){
	if (f.GetDimX() <= 0 && f.GetDimY() <= 0){
		return;
	}

	glm::vec2 v;
	int cntr = 0;
	glBegin(GL_LINES);
	for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dimx)){
		for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dimy)) {
			f.Sample(x, y, v);

			glVertex2f(x, y);
			glVertex2f(x+v[0]*mulConst, y+v[1]*mulConst);
			cntr++;

		}
	}
	glEnd();
}


void DrawVectorFieldNormalized(Field<glm::vec2> & f,  int dimx, int dimy, float mulConst = 1.0f){
	glm::vec2 v;
	glBegin(GL_LINES);
	for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dimx)){
		for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dimy)) {
			f.Sample(x, y, v);
			if (std::fabs(v[0]) < std::numeric_limits<float>::epsilon() &&
				std::fabs(v[1]) < std::numeric_limits<float>::epsilon()){
				v[0]=v[1] = 0;	
			}
			else {
				v = glm::normalize(v);
			}

			glVertex2f(x, y);
			glVertex2f(x+v[0]*mulConst, y+v[1]*mulConst);
		}
	}
	glEnd();
}

void DrawField() {
	field = fs->GetCurrField();

	glm::vec2 vel;
	float scal;
	int dims[2];
	dims[0] = field->GetDimX();
	dims[1] = field->GetDimY();

	glColor3f(0.0,1.0,0.0);
	DrawScalarField( *(fs->GetCurrScalarField() ),
		fs->GetCurrScalarField()->GetDimX(),
		fs->GetCurrScalarField()->GetDimY());


	// Draw Curl

	
	//glBegin(GL_POINTS);
	//for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dims[0])){
	//	for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dims[1])) {
	//		curl.Sample(x, y, scal);
	//		glColor3f(scal > 0 ? scal : 0,
	//				scal > 0 ? scal : 0,
	//				scal < 0 ? -scal : 0);
	//		glVertex2f(x+vel[0], y+vel[1]);

	//	}
	//}

	//glEnd();

	glPointSize(1.0f);

	// Velocity

	if (doDrawVectorField){
		glColor3f(1.0, 0.0, 0.0);
		DrawVectorField(*fs->GetCurrField(), fs->GetCurrField()->GetDimX(), fs->GetCurrField()->GetDimY(), 0.1f);
		
	}
	//glColor3f(0,1,0);
	//DrawVectorField(testField, testField.GetDimX(), testField.GetDimY(), 0.1f);

	//// Curl fied gradient
	

	//if (waveField.GetDimX() > 0){
	//	glColor3f(0.0,1.0,0.0);
	//	DrawVectorField(curlField, curlField.GetDimX(), curlField.GetDimY(), 0.1f);
	//	glColor3f(1.0,0.0,0.0);
	//	//DrawVectorField(waveField, curlField.GetDimX(), curlField.GetDimY(),0.1f);
	//}

	//if (doDecomp){

	//	//DrawVectorField(curlField, curlField.GetDimX(), curlField.GetDimY());
	//	glColor3f(1.0, 0.0, 0.0);
	//	//DrawScalarField(swaveField, swaveField.GetDimX(), swaveField.GetDimY());
	//	DrawVectorField(waveField, waveField.GetDimX(), waveField.GetDimY(),0.1f);

	//	//DrawVectorField(*fs->GetCurrField(), fs->GetCurrField()->GetDimX(), fs->GetCurrField()->GetDimY(), 0.1f);
	//}
	//else {
	//	

	//}
	
	//glColor3f(1.0,0.0,0.0);
	//DrawVectorField(curlField, curlField.GetDimX()*2, curlField.GetDimX()*2);

	//glBegin(GL_LINES);
	//for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dims[0])){
	//	for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dims[1])) {
	//		curlGrad.Sample(x, y, vel);
	//		glVertex2f(x, y);
	//		glVertex2f(x+vel[0], y+vel[1]);
	//	}
	//}

	//glEnd();
 //   
	//glBegin(GL_LINES);
	//for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dims[0])){
	//	for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dims[1])) {
	//		curlGrad.Sample(x, y, vel);
 //   
 //   
	//		glVertex2f(x, y);
	//		glVertex2f(x+vel[0]*1.0/fluidDimX, y+vel[1]*1.0/fluidDimX);
	//	}
 //   }
 //   
	//glEnd();
	/*glColor3f(0.0,1.0,0.0);
	glBegin(GL_LINES);
	for (float x = 0.0f; x < 1.0f; x += 1.0f/(float)(dims[0])){
		for ( float y = 0.0f; y < 1.0f; y += 1.0f/(float)(dims[1])) {
			curlField.Sample(x, y, vel);


			glVertex2f(x, y);
			glVertex2f(x+vel[0], y+vel[1]);
		}
	}

	glEnd();*/
}



void DrawVector(const glm::vec2 & inPos,const glm::vec2 & inVec, float mult = 1.0f){
	glBegin(GL_LINES);
	glVertex3fv(&inPos[0]);
	glVertex3f(inPos.x+inVec.x*mult, inPos.y+inVec.y*mult, 0);
	glEnd();


}

void DrawHaarCell( const HaarCell & inCell) {
	float cellW = inCell.cellW;
	int cellX = inCell.cellX;
	int cellY = inCell.cellY;
	
	glColor3f(0.0, 0.2, 0.5);
	glBegin(GL_QUADS);

	glVertex2f((float)cellX*cellW, (float)cellY*cellW);
	glVertex2f((float)cellX*cellW+cellW, (float)cellY*cellW);
	glVertex2f((float)cellX*cellW+cellW, (float)cellY*cellW+cellW);
	glVertex2f((float)cellX*cellW, (float)cellY*cellW+cellW);

	glEnd();

	glColor3f(0.0, 1.0, 0.0);
	glm::vec2 velPos( cellX*cellW+0.5f*cellW, cellY*cellW+0.5f*cellW);
	DrawVector(velPos, inCell.vel, 0.1f);

}

void DrawSelectedCell() {


	int cellX = haarSelector.wvCellX,
		cellY = haarSelector.wvCellY;
	float cellW = haarSelector.cellW;
	glColor3f(0.7, 0, 0.1);
	
	glBegin(GL_QUADS);

	glVertex2f((float)cellX*cellW, (float)cellY*cellW);
	glVertex2f((float)cellX*cellW+cellW, (float)cellY*cellW);
	glVertex2f((float)cellX*cellW+cellW, (float)cellY*cellW+cellW);
	glVertex2f((float)cellX*cellW, (float)cellY*cellW+cellW);

	glEnd();

	glColor3f(0.0, 1.0, 0.0);
	glm::vec2 velPos( cellX*cellW+0.5f*cellW, cellY*cellW+0.5f*cellW);
	DrawVector(velPos, haarSelector.vel, 0.1f);


	for (int i = 0; i < haarSelector.haarCellVec.size(); ++i){
		DrawHaarCell(haarSelector.haarCellVec[i]);
	}
}


//Display
//This function will be called as fast as possible, unlike the Update func.
//Only GL display code should be put in this section.  The screen
//is automatically cleared to white in the InitFrames() func, and the 
//glfwSwapBuffers() is called by the default EndFrames() function, so you only
//have to worry about actual drawing here.
void Display() {
	DrawField();

	if (doDrawSelectedCells){
		DrawSelectedCell();
	}
}

//Update
//This function will be called every timestep (ts).  The timestep can be
//set via the SetTimeStep method in SimulationManager. Put all non-drawing code
//here.


void convertScreenToFluidPos(
	int lastMouseX,
	int lastMouseY,
	float & cellX,
	float & cellY){
		
		cellX = ((float)lastMouseX / (float)WINDOW_W) * fluidDimX;
		cellY = ((float)lastMouseY / (float)WINDOW_H) * fluidDimY;
}

int lastMouseState = GLFW_RELEASE;



int lmx, lmy;
int ldx, ldy;
bool CheckMouse() {
	// Get position
	int mx, my;
	glfwGetMousePos(&mx, &my);

	float cellX, cellY;
	convertScreenToFluidPos(mx, my, cellX, cellY);
	haarSelector.SetSelectedCell(cellX, fluidDimY - cellY);
	
	bool retval;

	if ( glfwGetMouseButton( GLFW_MOUSE_BUTTON_LEFT ) ==
		GLFW_PRESS ){

		my = WINDOW_H - my;

		if (lastMouseState == GLFW_PRESS){
			ldx = mx - lmx;
			ldy = my - lmy;
			retval = true;
		}
		else {
			retval = false;
		}

		lmx = mx;
		lmy = my;
		lastMouseState = GLFW_PRESS;
	}
	else {
		lastMouseState = GLFW_RELEASE;
		retval = false;
	}

	if( glfwGetMouseButton( GLFW_MOUSE_BUTTON_RIGHT) == GLFW_PRESS ){
		float cellX, cellY;

		my = WINDOW_H - my;
		convertScreenToFluidPos(mx, my, cellX, cellY);

		Field<float> * sf = fs->GetCurrScalarField();

		for (int i = -skernel; i < skernel; ++i){
			for (int j = -skernel; j < skernel; ++j){
				if (cellX+i >0 && cellX +i < sf->GetDimX()-1 &&
					cellY+j>0 && cellY+j<sf->GetDimX()-1){

					(*sf)(cellX+i,cellY+j) = 1;
				}
			}
		}
	}

	return retval;
	

}


void MouseCallback( int x, int y ){

	Field<glm::vec2> * f = fs->GetCurrField();

	if (CheckMouse()){
		// Add force to grid.
		//float diff = ( (float)WINDOW_W - (1+(2*fieldBorder))*WINDOW_W)

		float cellX, cellY;
		convertScreenToFluidPos( lmx, lmy, cellX, cellY);

		glm::vec2 dirvec(ldx, ldy);
		dirvec = glm::normalize(dirvec);

		for (int i = -fkernel; i < fkernel; ++i){
			for (int j = -fkernel; j < fkernel; ++j){
				(*f)(cellX+i,cellY+j) += glm::vec2(dirvec.x*appliedForce, dirvec.y*appliedForce);
			}
		}
	}	
}

void KeyCallback( int key, int action ){
	if (action == GLFW_PRESS){
		switch(key){
		case 'W':
			if (key == 'W' && action == GLFW_PRESS){
				doDecomp = !doDecomp;
				printf("DoDecomp: %d\n", doDecomp);
			}
			break;
		case 'Q' :
			decompLev++;
			break;
		default:break;
		}
	} 
}

void ComputeCurl( const Field<glm::vec2> & velField,
	Field<float> & outCurl ){

	float dx = 1.0f/fluidDimX,
		dy = 1.0f/fluidDimY;

	float max = std::numeric_limits<float>::min();

	for (int i = 0; i < fluidDimX; ++i){
		for (int j = 0; j < fluidDimY; ++j){
			float c =
				(
					(velField(i+1,j)[1] - (velField(i-1,j)[1]))/dy
					-
					(velField(i,j+1)[0] - (velField(i,j-1)[0]))/dx
				);

			outCurl(i, j) = c;
			if (c > max)
			{
				max = c;
			}
		}
	}

	// Normalize
	//for (int i = 0; i < fluidDimX; ++i){
	//	for (int j = 0; j < fluidDimY; ++j){
	//		//curl(i,j)/= max;
	//		curl(i,j) = std::fabs(curl(i,j));
	//	}
	//}

}

void ComputeCurlGrad(Field<float> & curl, Field<glm::vec2> & curlGrad){

	float h = 1.0f/fluidDimX;
	for (int i = 0; i < fluidDimX; ++i){
		for (int j = 0; j < fluidDimY; ++j){
			float dwdx = (curl(i+1, j) - curl(i-1, j))*0.5f*h;
			float dwdy = (curl(i, j+1) - curl(i, j-1))*0.5f*h;
			curlGrad(i,j) = glm::vec2( dwdx, dwdy );
			curlField(i,j) = glm::vec2( dwdy, -dwdx );
			diffField(i,j) = (*fs->GetCurrField())(i,j) - curlField(i,j);
		}
	}

}

void ConfineVorticy(float energy){
	Field<glm::vec2> * currField = fs->GetCurrField();
	float h = 1.0f/fluidDimX;
	ComputeCurl(*fs->GetCurrField(), curl);
	ComputeCurlGrad(curl, curlGrad);

	glm::vec2 normGrad;
	
	for (int i = 0; i < fluidDimX; ++i){
		for (int j = 0; j < fluidDimY; ++j){
			normGrad = curlGrad(i,j);
			float dot = glm::dot(normGrad, normGrad);
			
			if (dot > std::numeric_limits<float>::epsilon()){
				dot = std::sqrt(dot);
				normGrad = normGrad/dot;
			}
			else {
				normGrad.x = 0;
				normGrad.y = 0;
			}

			float curlf = curl(i,j);


			(*currField)(i,j) += energy*glm::vec2(normGrad.y*curlf, -normGrad.x*curlf)*h; 
            //(*currField)(i,j) += energy*glm::vec2(curlField(i,j).x, curlField(i,j).y)*h;
			curlGrad(i,j) = glm::vec2(normGrad.y*curlf, -normGrad.x*curlf);
		}
	}

}

void Update(float ts) {
	//XXX
	//All update code goes here
	//...
	// Set
	CheckMouse();
	
	fs->SetSelector(&haarSelector);

	// Add voritcy confinement
	ConfineVorticy(vcEnergy);

	haarWv.SetData(*(fs->GetCurrField()->GetRawData()), fluidDimX);
	
	haarSelector.dimX = fluidDimX;
	haarSelector.dimY = fluidDimY;
	haarSelector.numLevels = haarWv.GetDecomp()->numLevs;

	haarWv.Decompose();
	HaarDecomp * decomp = haarWv.GetDecomp();
	//for (int c = 0; c < decomp->coeffDim/2; ++c){
	//	for (int r = 0; r < decomp->coeffDim/2; ++r){
	//		decomp->mData[r*decomp->totalDim + c] *= glm::vec2(0,1);
	//	}
	//}


	//haarWv.Decompose(*decomp);

	haarWv.Reconstruct();

	//
	//XXX
	//int testfieldim = decomp->coeffDim;
    int testfieldim = fluidDimX;
    //std::vector<glm::vec2> * data = &(decomp->mData);
	std::vector<glm::vec2> * data = haarWv.GetData();



	testField = Field<glm::vec2>(testfieldim, testfieldim);
	for (int i = 0; i < testfieldim; ++i){
		for (int j = 0; j < testfieldim; ++j){
			testField(i,j) = (*data)[j*fluidDimX+i];
		}
	}

	//fs->SetCurrField(testField);

	//haarWv.Reconstruct();

	//std::vector<glm::vec2> * hd = haarWv.GetData();
	//Field<glm::vec2> nf(*hd, fluidDimX, fluidDimX);
	//fs->SetCurrField(nf);
	
	// XXX
	// For the Samavati wavelets
	//Field<glm::vec2> & myfield = (*fs->GetCurrField());
	//waveField = myfield;

	//if (doDecomp){
	//	int lev = 0;
	//	while(lev < decompLev &&
	//		DecompField(waveField, waveField)
	//		){
	//		lev++;
	//	}
	//} 
	//else {
	//	decompLev = 0;
	//}
	//DecompField(swaveField, swaveField);

	//DecompWavelets(curlField, waveField);
	//DecompWavelets(waveField, waveField);
	//DecompField(waveField, waveField);
	//DecompField(waveField, waveField);

	fs->Update();

	//...
	//XXX	
}

void Init(){

	
	sm.SetTimeStep(ts);


	//XXX
	//Can initialize the camera and simulator params (callbacks) here
	//...
	fs = new FluidSolver(fluidDimX, fluidDimY, ts);
	fs->SetInterpMode(FluidSolver::BICUBIC);
	fs->SetViscosity(fluidVisc);


	curl = Field<float>(fluidDimX, fluidDimY);
	curlGrad = Field<glm::vec2>(fluidDimX, fluidDimY);
	curlField = Field<glm::vec2>(fluidDimX, fluidDimY);
	diffField = Field<glm::vec2>(fluidDimX, fluidDimY);

	// InputCallback
	glfwSetMousePosCallback(MouseCallback);
	glfwSetKeyCallback(KeyCallback);
	
}

void SetDoDrawVelocity() {
	doDrawVectorField = !doDrawVectorField;
}

void IncWaveLev(){
	haarSelector.IncLevel();
}

void DecWaveLev() {
	haarSelector.DecLevel();
}

void AddHaarCell(){
	HaarCell newCell;
	haarSelector.GetHaarCell(newCell);
	haarSelector.haarCellVec.push_back(newCell);
}

void RotateHaarVec() {
	glm::vec2 vel = haarSelector.vel;

	// Rotates clockwise
	float x = vel.x;
	vel.x = vel.y;
	vel.y = -x;

	haarSelector.vel = vel;
}

void DoDrawSelectedCells( ){
	doDrawSelectedCells = !doDrawSelectedCells;
}

int main (int argc, char ** argv ){


	sm.SetCallbackDisplay(Display);
	sm.SetCallbackUpdate(Update);
	sm.SetCallbackInitFrames(InitFrame);
	sm.SetCallbackLoadResources(Init);

	sm.AddKeyCallback('A', SetDoDrawVelocity);
	sm.AddKeyCallback(']', IncWaveLev);
	sm.AddKeyCallback('[', DecWaveLev);

	sm.AddKeyCallback('M', AddHaarCell);
	sm.AddKeyCallback('L', RotateHaarVec);
	sm.AddKeyCallback('S', DoDrawSelectedCells);

	//XXX

	//Start simulation.
	sm.Run();
}